import asyncio
import serial_asyncio
import serial
import threading
import time
import numpy as np
from transforms3d.euler import quat2euler, euler2quat
from ahrs.filters import Mahony
from pi_driver import I2cDriver
from .lpf import LowPassFilter
from .pid import PIDController


# control algorithm parameters
# velocity pid
pid_vel = PIDController(P=0.0137, I=0.03, D=0, ramp=10000, limit=np.pi / 9)
# stabilisation pid 40  100 1
pid_stb = PIDController(P=35, I=0.00001, D=1.4, ramp=100000, limit=12)
# velocity control filtering
lpf_pitch_cmd = LowPassFilter(Tf=0.07)
# low pass filters for user commands - throttle and steering
lpf_throttle = LowPassFilter(Tf=0.5)
lpf_steering = LowPassFilter(Tf=0.1)


class BalanceCar:
    def __init__(self, driver):
        self.max_throttle = 100 / 7.5
        self.max_steering = 100 / 40
        self.state = {
            "throttle": 0,
            "steering": 0,
            "running": False
        }
        self.sensor = I2cDriver()
        self.motor = driver
        self.active = False
        self.L = None
        self.R = None
        thread = threading.Thread(target=self.pid_loop)
        thread.daemon = True
        thread.start()

    def pid_loop(self):
        time_prev = time.time()
        dt = 0.01
        orientation = Mahony(frequency=int(1 / dt))
        Q_ = np.array(euler2quat(1.57, 0, 0))
        count = 0
        while True:
            if not self.active:
                time.sleep(0.3)
                continue
            now = time.time()
            if now - time_prev < dt:
                continue
            count = count + 1
            # print((now - time_prev)*1000)
            time_prev = now
            acc = np.array(self.sensor.readAccData(), dtype=np.float64)
            gyro = np.array(self.sensor.readGyroData(), dtype=np.float64)
            Q = orientation.updateIMU(
                Q_, gyr=gyro * 250 / 32768 / 180 * np.pi, acc=acc * 9.8 / 16384
            )
            Q_ = Q
            euler = np.array(quat2euler(Q))
            pitch, roll, yaw = euler
            pitch = pitch - 1.57
            pitch = -pitch
            if abs(pitch) > 1.2:
                self.set_torque(0, 0)
                self.state['running'] = False
                continue
            self.state['running'] = True
            target_pitch = lpf_pitch_cmd.calc(
                pid_vel.calc(
                    (self.motor.state[self.L].speed + (-self.motor.state[self.R].speed))
                    / 2
                    - lpf_throttle.calc(self.state["throttle"])
                )
            )
            # calculate the target voltage
            voltage_control = pid_stb.calc(target_pitch - pitch)
            steering_adj = lpf_steering.calc(self.state["steering"])
            self.set_torque(
                voltage_control + steering_adj, voltage_control - steering_adj
            )
            if count % 10 == 0:
                print(
                    self.motor.state[self.L].speed,
                    -self.motor.state[self.R].speed,
                    pitch,
                    target_pitch,
                    voltage_control,
                )
            # self.set_torque(-0.3, 0.2)

    def set_torque(self, t1, t2):
        self.motor.set_torque(self.L, t1)
        self.motor.set_torque(self.R, -t2)

    def set_motor(self, L, R):
        if L in self.motor.state and R in self.motor.state and L != R:
            self.L = L
            self.R = R

    def set_active(self, active):
        if active == True:
            if (
                self.L in self.motor.state
                and self.R in self.motor.state
                and self.L != self.R
            ):
                self.active = True
        else:
            if self.state['running'] == False:
                self.active = False

if __name__ == "__main__":
    car = BalanceCar("#BrushlessDriver")
    car.set_motor("B", "C")
    car.set_active(True)
    time.sleep(5)
